from langchain_community.llms import Ollama
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
import sys
import time
import requests
import json

def check_ollama_service():
    """
    检查Ollama服务是否正常运行

    返回值:
        bool: 服务可用返回True，否则返回False

    实现说明:
        通过发送GET请求到Ollama的API端点来检测服务状态
        当返回HTTP 200状态码时表示服务可用
    """
    try:
        response = requests.get("http://127.0.0.1:11434/api/tags")
        if response.status_code == 200:
            return True
    except Exception as e:
        print("Ollama 服务未运行,原因:", str(e))
        return False
    return False

def initialize_llm():
    """
    初始化LangChain的Ollama语言模型

    返回值:
        Ollama: 初始化成功的LLM模型实例

    异常处理:
        1. 如果Ollama服务不可用，程序将退出
        2. 如果模型初始化失败，程序将退出

    配置说明:
        - base_url: Ollama服务的本地地址
        - model: 使用的模型名称(deepseek-r1:14b)
    """
    if not check_ollama_service():
        print("错误: Ollama 服务未运行。请先启动 Ollama 服务。")
        print("启动命令: ollama serve")
        sys.exit(1)

    try:
        llm = Ollama(
            base_url="http://127.0.0.1:11434",
            model="deepseek-r1:14b"
        )
        return llm
    except Exception as e:
        print(f"初始化模型时出错: {str(e)}")
        sys.exit(1)

def stream_response(chain, input_data):
    """
    流式处理LLM响应并实时输出

    参数:
        chain (LLMChain): 已初始化的LLM链实例
        input_data (dict): 包含问题等输入数据的字典

    实现流程:
        1. 格式化提示词并显示问题
        2. 发送POST请求到Ollama的generate API
        3. 实时处理流式响应并输出
        4. 处理可能的错误情况

    技术要点:
        - 使用requests的stream=True参数实现流式传输
        - 通过iter_lines()逐行处理响应
        - 从JSON响应中提取并实时输出回答内容
    """
    print("\n问题:", chain.prompt.format(**input_data))
    print("\n回答:", end="", flush=True)

    try:
        prompt = chain.prompt.format(**input_data)
        response = requests.post(
            "http://127.0.0.1:11434/api/generate",
            json={
                "model": "deepseek-r1:14b",
                "prompt": prompt,
                "stream": True
            },
            stream=True,
            timeout=30
        )

        if response.status_code != 200:
            print(f"\nAPI 调用失败，状态码: {response.status_code}")
            return

        for line in response.iter_lines():
            if line:
                try:
                    json_response = json.loads(line)
                    if "response" in json_response:
                        # 只打印实际的回答内容
                        print(json_response["response"], end="", flush=True)
                except:
                    continue
        print("\n")
    except Exception as e:
        print(f"\n生成回答时出错: {str(e)}")

def main():
    """
    主程序入口，执行完整的LLM问答流程

    执行步骤:
        1. 初始化LLM模型
        2. 创建提示模板
        3. 构建LLM链
        4. 准备测试问题集
        5. 对每个问题执行流式问答

    设计说明:
        - 使用固定的提示模板确保回答格式一致
        - 问题列表涵盖不同领域测试模型能力
        - 流式输出提供更好的交互体验
        - 问题间添加1秒延迟避免输出混乱
    """
    # 初始化模型
    llm = initialize_llm()

    # 创建提示模板
    template = """你是一个有帮助的AI助手。请回答以下问题：
问题: {question}
回答:"""

    prompt = PromptTemplate(
        input_variables=["question"],
        template=template
    )

    # 创建 LLM 链
    chain = LLMChain(llm=llm, prompt=prompt)

    # 测试问题列表
    questions = [
        "请介绍一下你自己",
        "Python中的装饰器是什么？",
        "解释一下机器学习的原理",
        "如何提高代码质量？"
    ]

    # 流式输出每个问题的回答
    for question in questions:
        stream_response(chain, {"question": question})
        time.sleep(1)  # 添加短暂延迟，避免输出太快

if __name__ == "__main__":
    main() 